# -*- coding: utf-8 -*-
# @Time    : 2017/11/3 11:32
# @Author  : Deyu.Tian
# @Site    :
# @File    : config.py
# @Software: PyCharm Community Edition
from __future__ import print_function

import numpy as np
from landlab import RasterModelGrid
from landlab.components.overland_flow import OverlandFlow

import matplotlib.pyplot as plt

import config
import geotiff_tools as imt

def show_watermap(water_npy, imggt):
    """
    plot simulation results
    :param waterarray:
    :return:
    """

    water_arr = np.load(water_npy)
    #water_arr = water_arr.reshape((shape[0], shape[1]))
    print("max, min and shape of water-depth-grid: ", np.max(water_arr), np.min(water_arr), water_arr.shape)
    water_arr[water_arr > 0.001] = 1
    water_arr[water_arr < 0.001] = 0

    plt.figure(figsize=(30, 30), dpi=90)

    plt.subplot(131)
    plt.imshow(water_arr, cmap='Blues')
    plt.show()
    imt.array2rasterUTM("{}/result-waterdepth/waterDepth_muji_Binary.tif".format(config.basedir), imggt, water_arr)
    pass


def meltwater_flow(dem, source):
    """
    glacial meltwater flow simulations
    :param dem:
    :param source:
    :return:
    """
    run_time = 360  # duration of run, (s)
    alpha = 0.7  # time-step factor (nondimensional; from Bates et al., 2010)
    g = 9.81
    n = 0.5


    source = imt.img2array(source)
    source = source[:, 1:]
    shape = source.shape
    print("water-depth-grid shape:", source.shape)
    # source *= 0.0241  #init water depth m of kanswa
    # source *= 0.0265 #init water depth of upper gaiz
    source *= 0.0178 #init water depth of muji
    source = source + (1e-8)

    # #create a grid on which to calculate overland flow
    imggt, dem = imt.readmeltsnow(dem)
    # dem = dem[:, 1:]
    dem[dem == -9999] = 9999
    print("dem-grid shape:", dem.shape)
    x, y = dem.shape[0], dem.shape[1]
    grid = RasterModelGrid((x, y))
    grid.at_node['topographic__elevation'] = dem
    grid.at_node['surface_water__depth'] = source
    #grid.set_watershed_boundary_condition_outlet_id(6928, node_data=grid.at_node['topographic__elevation'], nodata_value=-9999)
    #grid.set_watershed_boundary_condition(grid.at_node['topographic__elevation'] ,nodata_value=-9999)

    # #init flow object and run simulation
    flow_obj = OverlandFlow(grid, steep_slopes=True, alpha=alpha, rainfall_intensity=0, g=g, mannings_n=n, theta=0.8)

    #flow_obj.run_one_step()
    elapsed_time = 0.0
    while elapsed_time < run_time:
        # First, we calculate our time step.
        dt = flow_obj.calc_time_step()
        # Now, we can generate overland flow.
        flow_obj.overland_flow()
        # Increased elapsed time
        # print('Elapsed time: ', elapsed_time)
        elapsed_time += dt


    #get the surface water depth outputs
    imt.array2rasterUTM("{}/result-waterdepth/waterDepth_muji.tif".format(config.basedir), imggt, grid.at_node['surface_water__depth'].reshape((shape[0], shape[1])))
    np.save("{}/result-waterdepth/waterDepth_muji.npy".format(config.basedir), grid.at_node['surface_water__depth'].reshape((shape[0], shape[1])))
    show_watermap("{}/result-waterdepth/waterDepth_muji.npy".format(config.basedir), imggt)
    pass


if __name__ == '__main__':
    # meltwater_flow("{}/DEMGrid_kanswa.tif".format(config.basedir), "{}/meltsnow_kanswa.tif".format(config.basedir))
    #imggt, dem = imt.readmeltsnow("{}/DEMGrid_kanswa.tif".format(config.basedir))
    #show_watermap("{}/result-waterdepth/waterDepth_kanswa.npy".format(config.basedir), imggt)
    # meltwater_flow("{}/DEMGrid_gaiz.tif".format(config.basedir), "{}/meltsnow_gaiz.tif".format(config.basedir))
    meltwater_flow("{}/DEMGrid_muji.tif".format(config.basedir), "{}/meltsnow_muji.tif".format(config.basedir))
    pass